class Solution
{
public:
    long long findMaximumElegance(vector<vector<int>> &items, int k)
    {
        sort(items.begin(), items.end(), [](const vector<int> &a, const vector<int> &b)
             { return a[0] > b[0]; });
        int n = items.size();
        unordered_map<int, int> categoryCount;
        priority_queue<int, vector<int>, greater<int>> pq;
        long long totalProfit = 0;

        for (int i = 0; i < k; ++i)
        {
            totalProfit += items[i][0];
            ++categoryCount[items[i][1]];
            if (categoryCount[items[i][1]] > 1)
            {
                pq.push(items[i][0]);
            }
        }

        long long maxElegance = totalProfit + categoryCount.size() * categoryCount.size();

        for (int i = k; i < n; ++i)
        {
            if (pq.empty())
            {
                break;
            }
            if (categoryCount.count(items[i][1]))
            {
                continue;
            }

            long long minProfit = pq.top();
            if (items[i][0] - minProfit + 2 * categoryCount.size() + 1 >= 0)
            {
                pq.pop();
                totalProfit -= minProfit;
                totalProfit += items[i][0];
                ++categoryCount[items[i][1]];
                maxElegance = max(maxElegance, totalProfit + static_cast<long long>(categoryCount.size() * categoryCount.size()));
            }
        }
        return maxElegance;
    }
};